"""
ros2订阅双目图像数据，分成左右两张发布，同时可按's'键保存
"""

import cv2
import numpy as np
import os
import time
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CompressedImage
from cv_bridge import CvBridge

class StereoCap(Node):
    def __init__(self):
        super().__init__('stereo_cap')
        self.cv_bridge = CvBridge()
        self.stereo_sub = self.create_subscription(CompressedImage, 
                                                   '/sensor_msgs/Image/Fisheye/compressed', 
                                                   self.stereo_callback, 1)
        self.left_pub = self.create_publisher(Image, '/left/image_raw', 1)
        self.right_pub = self.create_publisher(Image, '/right/image_raw', 1)

        self.img_count = 0
        self.save_path = "/home/daybeha/Documents/rosbag_tools/calibration/results"
        self.left_img_path = os.path.join(self.save_path, 'left/')
        self.right_img_path = os.path.join(self.save_path, 'right/')
        if not os.path.exists(self.left_img_path):
            os.makedirs(self.left_img_path)
        if not os.path.exists(self.right_img_path):
            os.makedirs(self.right_img_path)


        self.get_logger().info("Ready!")
        self.get_logger().info(f"Save path: {self.save_path}")
        self.get_logger().info("Press 's' to save image!")

    def stereo_callback(self, msg):
        # self.img = self.cv_bridge.compressed_imgmsg_to_cv2(msg, 'bgr8')
        # self.left_img = self.img[:, :int(self.img.shape[1]/2), :]
        # self.right_img = self.img[:, int(self.img.shape[1]/2):, :]

        self.img = self.cv_bridge.compressed_imgmsg_to_cv2(msg, 'passthrough')
        self.left_img = self.img[:, :int(self.img.shape[1]/2)]
        self.right_img = self.img[:, int(self.img.shape[1]/2):]

        cv2.imshow('left', self.left_img)
        cv2.imshow('right', self.right_img)
        if cv2.waitKey(1) & 0xFF == ord('s'):
            self.save_img()
            time.sleep(0.1)

        # left_msg = self.cv_bridge.cv2_to_imgmsg(self.left_img, 'bgr8')
        # right_msg = self.cv_bridge.cv2_to_imgmsg(self.right_img, 'bgr8')
        left_msg = self.cv_bridge.cv2_to_imgmsg(self.left_img, 'mono8')
        right_msg = self.cv_bridge.cv2_to_imgmsg(self.right_img, 'mono8')
        left_msg.header = msg.header
        right_msg.header = msg.header
        self.left_pub.publish(left_msg)
        self.right_pub.publish(right_msg)

    def save_img(self):
        if self.left_img is not None and self.right_img is not None:
            self.left_img_name = self.left_img_path + str(self.img_count) + '.jpg'
            self.right_img_name = self.right_img_path + str(self.img_count) + '.jpg'
            cv2.imwrite(self.left_img_name, self.left_img)
            cv2.imwrite(self.right_img_name, self.right_img)
            self.img_count += 1
            print(f"save image {self.img_count} success!")


def main(args=None):
    rclpy.init(args=args)
    stereo_cap = StereoCap()

    rclpy.spin(stereo_cap)
    stereo_cap.destroy_node()
    rclpy.shutdown()


if __name__ =="__main__":
    main()




